﻿import cv2
import numpy as np

# 参数设置
ALPHA = 1000  # 水印强度系数

# 1. 读取图像和水印
img = cv2.imread("input.png", cv2.IMREAD_GRAYSCALE)
watermark = cv2.imread("custom_watermark.png", cv2.IMREAD_GRAYSCALE)
assert img is not None and watermark is not None, "文件读取失败"

# 2. 调整水印尺寸
watermark = cv2.resize(watermark, (img.shape[1], img.shape[0]))

# 3. 转换为浮点型并归一化
img_float = img.astype(np.float32)
watermark_float = watermark.astype(np.float32) / 255.0

# 4. 扩展图像到最佳DFT尺寸
rows = cv2.getOptimalDFTSize(img.shape[0])
cols = cv2.getOptimalDFTSize(img.shape[1])
padded = cv2.copyMakeBorder(img_float, 0, rows-img.shape[0],
                           0, cols-img.shape[1], cv2.BORDER_CONSTANT)

# 5. 执行傅里叶变换
dft = cv2.dft(padded, flags=cv2.DFT_COMPLEX_OUTPUT)
dft_shift = np.fft.fftshift(dft)  # 中心化（可选）

# ===== 可视化原始频谱图 =====
magnitude_original = 20 * np.log(1 + cv2.magnitude(dft_shift[:,:,0], dft_shift[:,:,1]))
magnitude_original_normalized = cv2.normalize(magnitude_original, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
cv2.imwrite("original_spectrum.jpg", magnitude_original_normalized)


# 6. 嵌入水印到幅度谱
magnitude, phase = cv2.cartToPolar(dft_shift[:,:,0], dft_shift[:,:,1])

# 扩展水印尺寸
padded_watermark = cv2.copyMakeBorder(watermark_float, 0, rows-img.shape[0],
                                     0, cols-img.shape[1], cv2.BORDER_CONSTANT)
magnitude += ALPHA * padded_watermark  # 修改幅度谱

# ===== 可视化加水印后的频谱图 =====
magnitude_watermarked = 20 * np.log(1 + magnitude)
magnitude_watermarked_normalized = cv2.normalize(magnitude_watermarked, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
cv2.imwrite("watermarked_spectrum.jpg", magnitude_watermarked_normalized)

# 恢复复数表示
dft_shift[:,:,0], dft_shift[:,:,1] = cv2.polarToCart(magnitude, phase)

# 7. 逆傅里叶变换
dft_ishift = np.fft.ifftshift(dft_shift)  # 去中心化（如果之前做了中心化）
idft = cv2.idft(dft_ishift, flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT)

# 8. 裁剪和归一化
result = idft[:img.shape[0], :img.shape[1]]
result = cv2.normalize(result, None, 0, 255, cv2.NORM_MINMAX)
result = result.astype(np.uint8)

cv2.imwrite("output.jpg", result)